


<!DOCTYPE html>
<!--[if IE 8]><html class="no-js lt-ie9" lang="en" > <![endif]-->
<!--[if gt IE 8]><!--> <html class="no-js" lang="en" > <!--<![endif]-->
<head>
  <meta charset="utf-8">
  
  <meta name="viewport" content="width=device-width, initial-scale=1.0">
  
  <title>Automatic Mixed Precision examples &mdash; PyTorch master documentation</title>
  

  
  
  
  
    <link rel="canonical" href="https://pytorch.org/docs/stable/notes/amp_examples.html"/>
  

  

  
  
    

  

  <link rel="stylesheet" href="../_static/css/theme.css" type="text/css" />
  <!-- <link rel="stylesheet" href="../_static/pygments.css" type="text/css" /> -->
  <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.10.0-beta/dist/katex.min.css" type="text/css" />
  <link rel="stylesheet" href="../_static/css/jit.css" type="text/css" />
  <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.11.1/dist/katex.min.css" type="text/css" />
  <link rel="stylesheet" href="../_static/katex-math.css" type="text/css" />
    <link rel="index" title="Index" href="../genindex.html" />
    <link rel="search" title="Search" href="../search.html" />
    <link rel="next" title="Autograd mechanics" href="autograd.html" />
    <link rel="prev" title="PyTorch documentation" href="../index.html" /> 

  
  <script src="../_static/js/modernizr.min.js"></script>

  <!-- Preload the theme fonts -->

<link rel="preload" href="../_static/fonts/FreightSans/freight-sans-book.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="../_static/fonts/FreightSans/freight-sans-medium.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="../_static/fonts/IBMPlexMono/IBMPlexMono-Medium.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="../_static/fonts/FreightSans/freight-sans-bold.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="../_static/fonts/FreightSans/freight-sans-medium-italic.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="../_static/fonts/IBMPlexMono/IBMPlexMono-SemiBold.woff2" as="font" type="font/woff2" crossorigin="anonymous">

<!-- Preload the katex fonts -->

<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Math-Italic.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Main-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Main-Bold.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Size1-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Size4-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Size2-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Size3-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Caligraphic-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
</head>

<div class="container-fluid header-holder tutorials-header" id="header-holder">
  <div class="container">
    <div class="header-container">
      <a class="header-logo" href="https://pytorch.org/" aria-label="PyTorch"></a>

      <div class="main-menu">
        <ul>
          <li>
            <a href="https://pytorch.org/get-started">Get Started</a>
          </li>

          <li>
            <div class="ecosystem-dropdown">
              <a id="dropdownMenuButton" data-toggle="ecosystem-dropdown">
                Ecosystem
              </a>
              <div class="ecosystem-dropdown-menu">
                <a class="nav-dropdown-item" href="https://pytorch.org/hub"">
                  <span class=dropdown-title>Models (Beta)</span>
                  <p>Discover, publish, and reuse pre-trained models</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/ecosystem">
                  <span class=dropdown-title>Tools & Libraries</span>
                  <p>Explore the ecosystem of tools and libraries</p>
                </a>
              </div>
            </div>
          </li>

          <li>
            <a href="https://pytorch.org/mobile">Mobile</a>
          </li>

          <li>
            <a href="https://pytorch.org/blog/">Blog</a>
          </li>

          <li>
            <a href="https://pytorch.org/tutorials">Tutorials</a>
          </li>

          <li class="active">
            <a href="https://pytorch.org/docs/stable/index.html">Docs</a>
          </li>

          <li>
            <div class="resources-dropdown">
              <a id="resourcesDropdownButton" data-toggle="resources-dropdown">
                Resources
              </a>
              <div class="resources-dropdown-menu">
                <a class="nav-dropdown-item" href="https://pytorch.org/resources"">
                  <span class=dropdown-title>Developer Resources</span>
                  <p>Find resources and get questions answered</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/features">
                  <span class=dropdown-title>About</span>
                  <p>Learn about PyTorch’s features and capabilities</p>
                </a>
              </div>
            </div>
          </li>

          <li>
            <a href="https://github.com/pytorch/pytorch">Github</a>
          </li>
        </ul>
      </div>

      <a class="main-menu-open-button" href="#" data-behavior="open-mobile-menu"></a>
    </div>

  </div>
</div>


<body class="pytorch-body">

   

    

    <div class="table-of-contents-link-wrapper">
      <span>Table of Contents</span>
      <a href="#" class="toggle-table-of-contents" data-behavior="toggle-table-of-contents"></a>
    </div>

    <nav data-toggle="wy-nav-shift" class="pytorch-left-menu" id="pytorch-left-menu">
      <div class="pytorch-side-scroll">
        <div class="pytorch-menu pytorch-menu-vertical" data-spy="affix" role="navigation" aria-label="main navigation">
          <div class="pytorch-left-menu-search">
            

            
              
              
                <div class="version">
                  master (1.5.0 )
                </div>
              
            

            


  


<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../search.html" method="get">
    <input type="text" name="q" placeholder="Search Docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>

            
          </div>

          
<div>
  <a style="color:#F05732" href="https://pytorch.org/docs/stable/notes/amp_examples.html">
    You are viewing unstable developer preview docs.
    Click here to view docs for latest stable release.
  </a>
</div>

            
            
              
            
            
              <p class="caption"><span class="caption-text">Notes</span></p>
<ul class="current">
<li class="toctree-l1 current"><a class="current reference internal" href="#">Automatic Mixed Precision examples</a></li>
<li class="toctree-l1"><a class="reference internal" href="autograd.html">Autograd mechanics</a></li>
<li class="toctree-l1"><a class="reference internal" href="broadcasting.html">Broadcasting semantics</a></li>
<li class="toctree-l1"><a class="reference internal" href="cpu_threading_torchscript_inference.html">CPU threading and TorchScript inference</a></li>
<li class="toctree-l1"><a class="reference internal" href="cuda.html">CUDA semantics</a></li>
<li class="toctree-l1"><a class="reference internal" href="ddp.html">Distributed Data Parallel</a></li>
<li class="toctree-l1"><a class="reference internal" href="extending.html">Extending PyTorch</a></li>
<li class="toctree-l1"><a class="reference internal" href="faq.html">Frequently Asked Questions</a></li>
<li class="toctree-l1"><a class="reference internal" href="large_scale_deployments.html">Features for large-scale deployments</a></li>
<li class="toctree-l1"><a class="reference internal" href="multiprocessing.html">Multiprocessing best practices</a></li>
<li class="toctree-l1"><a class="reference internal" href="randomness.html">Reproducibility</a></li>
<li class="toctree-l1"><a class="reference internal" href="serialization.html">Serialization semantics</a></li>
<li class="toctree-l1"><a class="reference internal" href="windows.html">Windows FAQ</a></li>
</ul>
<p class="caption"><span class="caption-text">Language Bindings</span></p>
<ul>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/cppdocs/">C++ API</a></li>
<li class="toctree-l1"><a class="reference internal" href="../packages.html">Javadoc</a></li>
</ul>
<p class="caption"><span class="caption-text">Python API</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../torch.html">torch</a></li>
<li class="toctree-l1"><a class="reference internal" href="../nn.html">torch.nn</a></li>
<li class="toctree-l1"><a class="reference internal" href="../nn.functional.html">torch.nn.functional</a></li>
<li class="toctree-l1"><a class="reference internal" href="../tensors.html">torch.Tensor</a></li>
<li class="toctree-l1"><a class="reference internal" href="../tensor_attributes.html">Tensor Attributes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../tensor_view.html">Tensor Views</a></li>
<li class="toctree-l1"><a class="reference internal" href="../autograd.html">torch.autograd</a></li>
<li class="toctree-l1"><a class="reference internal" href="../cuda.html">torch.cuda</a></li>
<li class="toctree-l1"><a class="reference internal" href="../amp.html">torch.cuda.amp</a></li>
<li class="toctree-l1"><a class="reference internal" href="../distributed.html">torch.distributed</a></li>
<li class="toctree-l1"><a class="reference internal" href="../distributions.html">torch.distributions</a></li>
<li class="toctree-l1"><a class="reference internal" href="../hub.html">torch.hub</a></li>
<li class="toctree-l1"><a class="reference internal" href="../jit.html">torch.jit</a></li>
<li class="toctree-l1"><a class="reference internal" href="../nn.init.html">torch.nn.init</a></li>
<li class="toctree-l1"><a class="reference internal" href="../onnx.html">torch.onnx</a></li>
<li class="toctree-l1"><a class="reference internal" href="../optim.html">torch.optim</a></li>
<li class="toctree-l1"><a class="reference internal" href="../quantization.html">Quantization</a></li>
<li class="toctree-l1"><a class="reference internal" href="../rpc/index.html">Distributed RPC Framework</a></li>
<li class="toctree-l1"><a class="reference internal" href="../random.html">torch.random</a></li>
<li class="toctree-l1"><a class="reference internal" href="../sparse.html">torch.sparse</a></li>
<li class="toctree-l1"><a class="reference internal" href="../storage.html">torch.Storage</a></li>
<li class="toctree-l1"><a class="reference internal" href="../bottleneck.html">torch.utils.bottleneck</a></li>
<li class="toctree-l1"><a class="reference internal" href="../checkpoint.html">torch.utils.checkpoint</a></li>
<li class="toctree-l1"><a class="reference internal" href="../cpp_extension.html">torch.utils.cpp_extension</a></li>
<li class="toctree-l1"><a class="reference internal" href="../data.html">torch.utils.data</a></li>
<li class="toctree-l1"><a class="reference internal" href="../dlpack.html">torch.utils.dlpack</a></li>
<li class="toctree-l1"><a class="reference internal" href="../model_zoo.html">torch.utils.model_zoo</a></li>
<li class="toctree-l1"><a class="reference internal" href="../tensorboard.html">torch.utils.tensorboard</a></li>
<li class="toctree-l1"><a class="reference internal" href="../type_info.html">Type Info</a></li>
<li class="toctree-l1"><a class="reference internal" href="../named_tensor.html">Named Tensors</a></li>
<li class="toctree-l1"><a class="reference internal" href="../name_inference.html">Named Tensors operator coverage</a></li>
<li class="toctree-l1"><a class="reference internal" href="../__config__.html">torch.__config__</a></li>
</ul>
<p class="caption"><span class="caption-text">Libraries</span></p>
<ul>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/audio">torchaudio</a></li>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/text">torchtext</a></li>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/elastic/">TorchElastic</a></li>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/serve">TorchServe</a></li>
<li class="toctree-l1"><a class="reference external" href="http://pytorch.org/xla/">PyTorch on XLA Devices</a></li>
</ul>
<p class="caption"><span class="caption-text">Community</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../community/contribution_guide.html">PyTorch Contribution Guide</a></li>
<li class="toctree-l1"><a class="reference internal" href="../community/governance.html">PyTorch Governance</a></li>
<li class="toctree-l1"><a class="reference internal" href="../community/persons_of_interest.html">PyTorch Governance | Persons of Interest</a></li>
</ul>

            
          

        </div>
      </div>
    </nav>

    <div class="pytorch-container">
      <div class="pytorch-page-level-bar" id="pytorch-page-level-bar">
        <div class="pytorch-breadcrumbs-wrapper">
          















<div role="navigation" aria-label="breadcrumbs navigation">

  <ul class="pytorch-breadcrumbs">
    
      <li>
        <a href="../index.html">
          
            Docs
          
        </a> &gt;
      </li>

        
      <li>Automatic Mixed Precision examples</li>
    
    
      <li class="pytorch-breadcrumbs-aside">
        
            
            <a href="../_sources/notes/amp_examples.rst.txt" rel="nofollow"><img src="../_static/images/view-page-source-icon.svg"></a>
          
        
      </li>
    
  </ul>

  
</div>
        </div>

        <div class="pytorch-shortcuts-wrapper" id="pytorch-shortcuts-wrapper">
          Shortcuts
        </div>
      </div>

      <section data-toggle="wy-nav-shift" id="pytorch-content-wrap" class="pytorch-content-wrap">
        <div class="pytorch-content-left">

        
          
          <div class="rst-content">
          
            <div role="main" class="main-content" itemscope="itemscope" itemtype="http://schema.org/Article">
             <article itemprop="articleBody" id="pytorch-article" class="pytorch-article">
              
  <div class="section" id="automatic-mixed-precision-examples">
<span id="amp-examples"></span><h1>Automatic Mixed Precision examples<a class="headerlink" href="#automatic-mixed-precision-examples" title="Permalink to this headline">¶</a></h1>
<div class="admonition warning">
<p class="admonition-title">Warning</p>
<p><a class="reference internal" href="../amp.html#torch.cuda.amp.GradScaler" title="torch.cuda.amp.GradScaler"><code class="xref py py-class docutils literal notranslate"><span class="pre">torch.cuda.amp.GradScaler</span></code></a> is not a complete implementation of automatic mixed precision.
<a class="reference internal" href="../amp.html#torch.cuda.amp.GradScaler" title="torch.cuda.amp.GradScaler"><code class="xref py py-class docutils literal notranslate"><span class="pre">GradScaler</span></code></a> is only useful if you manually run regions of your model in <code class="docutils literal notranslate"><span class="pre">float16</span></code>.
If you aren’t sure how to choose op precision manually, the master branch and nightly pip/conda
builds include a context manager that chooses op precision automatically wherever it’s enabled.
See the <a class="reference external" href="https://pytorch.org/docs/master/amp.html">master documentation</a> for details.</p>
</div>
<div class="contents local topic" id="contents">
<ul class="simple">
<li><p><a class="reference internal" href="#gradient-scaling" id="id2">Gradient Scaling</a></p>
<ul>
<li><p><a class="reference internal" href="#typical-use" id="id3">Typical Use</a></p></li>
<li><p><a class="reference internal" href="#working-with-unscaled-gradients" id="id4">Working with Unscaled Gradients</a></p>
<ul>
<li><p><a class="reference internal" href="#gradient-clipping" id="id5">Gradient clipping</a></p></li>
</ul>
</li>
<li><p><a class="reference internal" href="#working-with-scaled-gradients" id="id6">Working with Scaled Gradients</a></p>
<ul>
<li><p><a class="reference internal" href="#gradient-penalty" id="id7">Gradient penalty</a></p></li>
</ul>
</li>
<li><p><a class="reference internal" href="#working-with-multiple-losses-and-optimizers" id="id8">Working with Multiple Losses and Optimizers</a></p></li>
</ul>
</li>
</ul>
</div>
<div class="section" id="gradient-scaling">
<span id="gradient-scaling-examples"></span><h2><a class="toc-backref" href="#id2">Gradient Scaling</a><a class="headerlink" href="#gradient-scaling" title="Permalink to this headline">¶</a></h2>
<p>Gradient scaling helps prevent gradient underflow when training with mixed precision,
as explained <a class="reference internal" href="../amp.html#gradient-scaling"><span class="std std-ref">here</span></a>.</p>
<p>Instances of <a class="reference internal" href="../amp.html#torch.cuda.amp.GradScaler" title="torch.cuda.amp.GradScaler"><code class="xref py py-class docutils literal notranslate"><span class="pre">torch.cuda.amp.GradScaler</span></code></a> help perform the steps of
gradient scaling conveniently, as shown in the following code snippets.</p>
<div class="section" id="typical-use">
<h3><a class="toc-backref" href="#id3">Typical Use</a><a class="headerlink" href="#typical-use" title="Permalink to this headline">¶</a></h3>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="c1"># Creates a GradScaler once at the beginning of training.</span>
<span class="n">scaler</span> <span class="o">=</span> <span class="n">GradScaler</span><span class="p">()</span>

<span class="k">for</span> <span class="n">epoch</span> <span class="ow">in</span> <span class="n">epochs</span><span class="p">:</span>
    <span class="k">for</span> <span class="nb">input</span><span class="p">,</span> <span class="n">target</span> <span class="ow">in</span> <span class="n">data</span><span class="p">:</span>
        <span class="n">optimizer</span><span class="o">.</span><span class="n">zero_grad</span><span class="p">()</span>
        <span class="n">output</span> <span class="o">=</span> <span class="n">model</span><span class="p">(</span><span class="nb">input</span><span class="p">)</span>
        <span class="n">loss</span> <span class="o">=</span> <span class="n">loss_fn</span><span class="p">(</span><span class="n">output</span><span class="p">,</span> <span class="n">target</span><span class="p">)</span>

        <span class="c1"># Scales the loss, and calls backward() on the scaled loss to create scaled gradients.</span>
        <span class="n">scaler</span><span class="o">.</span><span class="n">scale</span><span class="p">(</span><span class="n">loss</span><span class="p">)</span><span class="o">.</span><span class="n">backward</span><span class="p">()</span>

        <span class="c1"># scaler.step() first unscales the gradients of the optimizer&#39;s assigned params.</span>
        <span class="c1"># If these gradients do not contain infs or NaNs, optimizer.step() is then called,</span>
        <span class="c1"># otherwise, optimizer.step() is skipped.</span>
        <span class="n">scaler</span><span class="o">.</span><span class="n">step</span><span class="p">(</span><span class="n">optimizer</span><span class="p">)</span>

        <span class="c1"># Updates the scale for next iteration.</span>
        <span class="n">scaler</span><span class="o">.</span><span class="n">update</span><span class="p">()</span>
</pre></div>
</div>
</div>
<div class="section" id="working-with-unscaled-gradients">
<span id="id1"></span><h3><a class="toc-backref" href="#id4">Working with Unscaled Gradients</a><a class="headerlink" href="#working-with-unscaled-gradients" title="Permalink to this headline">¶</a></h3>
<p>All gradients produced by <code class="docutils literal notranslate"><span class="pre">scaler.scale(loss).backward()</span></code> are scaled.  If you wish to modify or inspect
the parameters’ <code class="docutils literal notranslate"><span class="pre">.grad</span></code> attributes between <code class="docutils literal notranslate"><span class="pre">backward()</span></code> and <code class="docutils literal notranslate"><span class="pre">scaler.step(optimizer)</span></code>,  you should
unscale them first.  For example, gradient clipping manipulates a set of gradients such that their global norm
(see <a class="reference internal" href="../nn.html#torch.nn.utils.clip_grad_norm_" title="torch.nn.utils.clip_grad_norm_"><code class="xref py py-func docutils literal notranslate"><span class="pre">torch.nn.utils.clip_grad_norm_()</span></code></a>) or maximum magnitude (see <a class="reference internal" href="../nn.html#torch.nn.utils.clip_grad_value_" title="torch.nn.utils.clip_grad_value_"><code class="xref py py-func docutils literal notranslate"><span class="pre">torch.nn.utils.clip_grad_value_()</span></code></a>)
is <span class="math"><span class="katex"><span class="katex-mathml"><math xmlns="http://www.w3.org/1998/Math/MathML"><semantics><mrow><mo>&lt;</mo><mo>=</mo></mrow><annotation encoding="application/x-tex">&lt;=</annotation></semantics></math></span><span class="katex-html" aria-hidden="true"><span class="base"><span class="strut" style="height:0.5782em;vertical-align:-0.0391em;"></span><span class="mrel">&lt;</span></span><span class="base"><span class="strut" style="height:0.36687em;vertical-align:0em;"></span><span class="mrel">=</span></span></span></span>

</span> some user-imposed threshold.  If you attempted to clip <em>without</em> unscaling, the gradients’ norm/maximum
magnitude would also be scaled, so your requested threshold (which was meant to be the threshold for <em>unscaled</em>
gradients) would be invalid.</p>
<p><code class="docutils literal notranslate"><span class="pre">scaler.unscale_(optimizer)</span></code> unscales gradients held by <code class="docutils literal notranslate"><span class="pre">optimizer</span></code>’s assigned parameters.
If your model or models contain other parameters that were assigned to another optimizer
(say <code class="docutils literal notranslate"><span class="pre">optimizer2</span></code>), you may call <code class="docutils literal notranslate"><span class="pre">scaler.unscale_(optimizer2)</span></code> separately to unscale those
parameters’ gradients as well.</p>
<div class="section" id="gradient-clipping">
<h4><a class="toc-backref" href="#id5">Gradient clipping</a><a class="headerlink" href="#gradient-clipping" title="Permalink to this headline">¶</a></h4>
<p>Calling <code class="docutils literal notranslate"><span class="pre">scaler.unscale_(optimizer)</span></code> before clipping enables you to clip unscaled gradients as usual:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">scaler</span> <span class="o">=</span> <span class="n">GradScaler</span><span class="p">()</span>

<span class="k">for</span> <span class="n">epoch</span> <span class="ow">in</span> <span class="n">epochs</span><span class="p">:</span>
    <span class="k">for</span> <span class="nb">input</span><span class="p">,</span> <span class="n">target</span> <span class="ow">in</span> <span class="n">data</span><span class="p">:</span>
        <span class="n">optimizer</span><span class="o">.</span><span class="n">zero_grad</span><span class="p">()</span>
        <span class="n">output</span> <span class="o">=</span> <span class="n">model</span><span class="p">(</span><span class="nb">input</span><span class="p">)</span>
        <span class="n">loss</span> <span class="o">=</span> <span class="n">loss_fn</span><span class="p">(</span><span class="n">output</span><span class="p">,</span> <span class="n">target</span><span class="p">)</span>
        <span class="n">scaler</span><span class="o">.</span><span class="n">scale</span><span class="p">(</span><span class="n">loss</span><span class="p">)</span><span class="o">.</span><span class="n">backward</span><span class="p">()</span>

        <span class="c1"># Unscales the gradients of optimizer&#39;s assigned params in-place</span>
        <span class="n">scaler</span><span class="o">.</span><span class="n">unscale_</span><span class="p">(</span><span class="n">optimizer</span><span class="p">)</span>

        <span class="c1"># Since the gradients of optimizer&#39;s assigned params are unscaled, clips as usual:</span>
        <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">utils</span><span class="o">.</span><span class="n">clip_grad_norm_</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">parameters</span><span class="p">(),</span> <span class="n">max_norm</span><span class="p">)</span>

        <span class="c1"># optimizer&#39;s gradients are already unscaled, so scaler.step does not unscale them,</span>
        <span class="c1"># although it still skips optimizer.step() if the gradients contain infs or NaNs.</span>
        <span class="n">scaler</span><span class="o">.</span><span class="n">step</span><span class="p">(</span><span class="n">optimizer</span><span class="p">)</span>

        <span class="c1"># Updates the scale for next iteration.</span>
        <span class="n">scaler</span><span class="o">.</span><span class="n">update</span><span class="p">()</span>
</pre></div>
</div>
<p><code class="docutils literal notranslate"><span class="pre">scaler</span></code> records that <code class="docutils literal notranslate"><span class="pre">scaler.unscale_(optimizer)</span></code> was already called for this optimizer
this iteration, so <code class="docutils literal notranslate"><span class="pre">scaler.step(optimizer)</span></code> knows not to redundantly unscale gradients before
(internally) calling <code class="docutils literal notranslate"><span class="pre">optimizer.step()</span></code>.</p>
<div class="admonition warning">
<p class="admonition-title">Warning</p>
<p><code class="xref py py-meth docutils literal notranslate"><span class="pre">unscale_()</span></code> should only be called once per optimizer per <code class="xref py py-meth docutils literal notranslate"><span class="pre">step()</span></code> call,
and only after all gradients for that optimizer’s assigned parameters have been accumulated.
Calling <code class="xref py py-meth docutils literal notranslate"><span class="pre">unscale_()</span></code> twice for a given optimizer between each <code class="xref py py-meth docutils literal notranslate"><span class="pre">step()</span></code> triggers a RuntimeError.</p>
</div>
</div>
</div>
<div class="section" id="working-with-scaled-gradients">
<h3><a class="toc-backref" href="#id6">Working with Scaled Gradients</a><a class="headerlink" href="#working-with-scaled-gradients" title="Permalink to this headline">¶</a></h3>
<p>For some operations, you may need to work with scaled gradients in a setting where
<code class="docutils literal notranslate"><span class="pre">scaler.unscale_</span></code> is unsuitable.</p>
<div class="section" id="gradient-penalty">
<h4><a class="toc-backref" href="#id7">Gradient penalty</a><a class="headerlink" href="#gradient-penalty" title="Permalink to this headline">¶</a></h4>
<p>A gradient penalty implementation typically creates gradients out-of-place using
<a class="reference internal" href="../autograd.html#torch.autograd.grad" title="torch.autograd.grad"><code class="xref py py-func docutils literal notranslate"><span class="pre">torch.autograd.grad()</span></code></a>, combines them to create the penalty value,
and adds the penalty value to the loss.</p>
<p>Here’s an ordinary example of an L2 penalty without gradient scaling:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="k">for</span> <span class="n">epoch</span> <span class="ow">in</span> <span class="n">epochs</span><span class="p">:</span>
    <span class="k">for</span> <span class="nb">input</span><span class="p">,</span> <span class="n">target</span> <span class="ow">in</span> <span class="n">data</span><span class="p">:</span>
        <span class="n">optimizer</span><span class="o">.</span><span class="n">zero_grad</span><span class="p">()</span>
        <span class="n">output</span> <span class="o">=</span> <span class="n">model</span><span class="p">(</span><span class="nb">input</span><span class="p">)</span>
        <span class="n">loss</span> <span class="o">=</span> <span class="n">loss_fn</span><span class="p">(</span><span class="n">output</span><span class="p">,</span> <span class="n">target</span><span class="p">)</span>

        <span class="c1"># Creates some gradients out-of-place</span>
        <span class="n">grad_params</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">autograd</span><span class="o">.</span><span class="n">grad</span><span class="p">(</span><span class="n">loss</span><span class="p">,</span> <span class="n">model</span><span class="o">.</span><span class="n">parameters</span><span class="p">(),</span> <span class="n">create_graph</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>

        <span class="c1"># Computes the penalty term and adds it to the loss</span>
        <span class="n">grad_norm</span> <span class="o">=</span> <span class="mi">0</span>
        <span class="k">for</span> <span class="n">grad</span> <span class="ow">in</span> <span class="n">grad_params</span><span class="p">:</span>
            <span class="n">grad_norm</span> <span class="o">+=</span> <span class="n">grad</span><span class="o">.</span><span class="n">pow</span><span class="p">(</span><span class="mi">2</span><span class="p">)</span><span class="o">.</span><span class="n">sum</span><span class="p">()</span>
        <span class="n">grad_norm</span> <span class="o">=</span> <span class="n">grad_norm</span><span class="o">.</span><span class="n">sqrt</span><span class="p">()</span>
        <span class="n">loss</span> <span class="o">=</span> <span class="n">loss</span> <span class="o">+</span> <span class="n">grad_norm</span>

        <span class="n">loss</span><span class="o">.</span><span class="n">backward</span><span class="p">()</span>
        <span class="n">optimizer</span><span class="o">.</span><span class="n">step</span><span class="p">()</span>
</pre></div>
</div>
<p>To implement a gradient penalty <em>with</em> gradient scaling, the loss passed to
<a class="reference internal" href="../autograd.html#torch.autograd.grad" title="torch.autograd.grad"><code class="xref py py-func docutils literal notranslate"><span class="pre">torch.autograd.grad()</span></code></a> should be scaled.  The resulting out-of-place gradients
will therefore be scaled, and should be unscaled before being combined to create the
penalty value.</p>
<p>Here’s how that looks for the same L2 penalty:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">scaler</span> <span class="o">=</span> <span class="n">GradScaler</span><span class="p">()</span>

<span class="k">for</span> <span class="n">epoch</span> <span class="ow">in</span> <span class="n">epochs</span><span class="p">:</span>
    <span class="k">for</span> <span class="nb">input</span><span class="p">,</span> <span class="n">target</span> <span class="ow">in</span> <span class="n">data</span><span class="p">:</span>
        <span class="n">optimizer</span><span class="o">.</span><span class="n">zero_grad</span><span class="p">()</span>
        <span class="n">output</span> <span class="o">=</span> <span class="n">model</span><span class="p">(</span><span class="nb">input</span><span class="p">)</span>
        <span class="n">loss</span> <span class="o">=</span> <span class="n">loss_fn</span><span class="p">(</span><span class="n">output</span><span class="p">,</span> <span class="n">target</span><span class="p">)</span>

        <span class="c1"># Scales the loss for the out-of-place backward pass, resulting in scaled grad_params</span>
        <span class="n">scaled_grad_params</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">autograd</span><span class="o">.</span><span class="n">grad</span><span class="p">(</span><span class="n">scaler</span><span class="o">.</span><span class="n">scale</span><span class="p">(</span><span class="n">loss</span><span class="p">),</span> <span class="n">model</span><span class="o">.</span><span class="n">parameters</span><span class="p">(),</span> <span class="n">create_graph</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>

        <span class="c1"># Unscales grad_params before computing the penalty.  grad_params are not owned</span>
        <span class="c1"># by any optimizer, so ordinary division is used instead of scaler.unscale_:</span>
        <span class="n">inv_scale</span> <span class="o">=</span> <span class="mf">1.</span><span class="o">/</span><span class="n">scaler</span><span class="o">.</span><span class="n">get_scale</span><span class="p">()</span>
        <span class="n">grad_params</span> <span class="o">=</span> <span class="p">[</span><span class="n">p</span><span class="o">*</span><span class="n">inv_scale</span> <span class="k">for</span> <span class="n">p</span> <span class="ow">in</span> <span class="n">scaled_grad_params</span><span class="p">]</span>

        <span class="c1"># Computes the penalty term and adds it to the loss</span>
        <span class="n">grad_norm</span> <span class="o">=</span> <span class="mi">0</span>
        <span class="k">for</span> <span class="n">grad</span> <span class="ow">in</span> <span class="n">grad_params</span><span class="p">:</span>
            <span class="n">grad_norm</span> <span class="o">+=</span> <span class="n">grad</span><span class="o">.</span><span class="n">pow</span><span class="p">(</span><span class="mi">2</span><span class="p">)</span><span class="o">.</span><span class="n">sum</span><span class="p">()</span>
        <span class="n">grad_norm</span> <span class="o">=</span> <span class="n">grad_norm</span><span class="o">.</span><span class="n">sqrt</span><span class="p">()</span>
        <span class="n">loss</span> <span class="o">=</span> <span class="n">loss</span> <span class="o">+</span> <span class="n">grad_norm</span>

        <span class="c1"># Applies scaling to the backward call as usual.  Accumulates leaf gradients that are correctly scaled.</span>
        <span class="n">scaler</span><span class="o">.</span><span class="n">scale</span><span class="p">(</span><span class="n">loss</span><span class="p">)</span><span class="o">.</span><span class="n">backward</span><span class="p">()</span>

        <span class="c1"># step() and update() proceed as usual.</span>
        <span class="n">scaler</span><span class="o">.</span><span class="n">step</span><span class="p">(</span><span class="n">optimizer</span><span class="p">)</span>
        <span class="n">scaler</span><span class="o">.</span><span class="n">update</span><span class="p">()</span>
</pre></div>
</div>
</div>
</div>
<div class="section" id="working-with-multiple-losses-and-optimizers">
<h3><a class="toc-backref" href="#id8">Working with Multiple Losses and Optimizers</a><a class="headerlink" href="#working-with-multiple-losses-and-optimizers" title="Permalink to this headline">¶</a></h3>
<p>If your network has multiple losses, you must call <code class="docutils literal notranslate"><span class="pre">scaler.scale</span></code> on each of them individually.
If your network has multiple optimizers, you may call <code class="docutils literal notranslate"><span class="pre">scaler.unscale_</span></code> on any of them individually,
and you must call <code class="docutils literal notranslate"><span class="pre">scaler.step</span></code> on each of them individually.</p>
<p>However, <code class="docutils literal notranslate"><span class="pre">scaler.update()</span></code> should only be called once,
after all optimizers used this iteration have been stepped:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">scaler</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cuda</span><span class="o">.</span><span class="n">amp</span><span class="o">.</span><span class="n">GradScaler</span><span class="p">()</span>

<span class="k">for</span> <span class="n">epoch</span> <span class="ow">in</span> <span class="n">epochs</span><span class="p">:</span>
    <span class="k">for</span> <span class="nb">input</span><span class="p">,</span> <span class="n">target</span> <span class="ow">in</span> <span class="n">data</span><span class="p">:</span>
        <span class="n">optimizer0</span><span class="o">.</span><span class="n">zero_grad</span><span class="p">()</span>
        <span class="n">optimizer1</span><span class="o">.</span><span class="n">zero_grad</span><span class="p">()</span>
        <span class="n">output0</span> <span class="o">=</span> <span class="n">model0</span><span class="p">(</span><span class="nb">input</span><span class="p">)</span>
        <span class="n">output1</span> <span class="o">=</span> <span class="n">model1</span><span class="p">(</span><span class="nb">input</span><span class="p">)</span>
        <span class="n">loss0</span> <span class="o">=</span> <span class="n">loss_fn</span><span class="p">(</span><span class="mi">2</span> <span class="o">*</span> <span class="n">output0</span> <span class="o">+</span> <span class="mi">3</span> <span class="o">*</span> <span class="n">output1</span><span class="p">,</span> <span class="n">target</span><span class="p">)</span>
        <span class="n">loss1</span> <span class="o">=</span> <span class="n">loss_fn</span><span class="p">(</span><span class="mi">3</span> <span class="o">*</span> <span class="n">output0</span> <span class="o">-</span> <span class="mi">5</span> <span class="o">*</span> <span class="n">output1</span><span class="p">,</span> <span class="n">target</span><span class="p">)</span>

        <span class="n">scaler</span><span class="o">.</span><span class="n">scale</span><span class="p">(</span><span class="n">loss0</span><span class="p">)</span><span class="o">.</span><span class="n">backward</span><span class="p">(</span><span class="n">retain_graph</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
        <span class="n">scaler</span><span class="o">.</span><span class="n">scale</span><span class="p">(</span><span class="n">loss1</span><span class="p">)</span><span class="o">.</span><span class="n">backward</span><span class="p">()</span>

        <span class="c1"># You can choose which optimizers receive explicit unscaling, if you</span>
        <span class="c1"># want to inspect or modify the gradients of the params they own.</span>
        <span class="n">scaler</span><span class="o">.</span><span class="n">unscale_</span><span class="p">(</span><span class="n">optimizer0</span><span class="p">)</span>

        <span class="n">scaler</span><span class="o">.</span><span class="n">step</span><span class="p">(</span><span class="n">optimizer0</span><span class="p">)</span>
        <span class="n">scaler</span><span class="o">.</span><span class="n">step</span><span class="p">(</span><span class="n">optimizer1</span><span class="p">)</span>

        <span class="n">scaler</span><span class="o">.</span><span class="n">update</span><span class="p">()</span>
</pre></div>
</div>
<p>Each optimizer independently checks its gradients for infs/NaNs, and therefore makes an independent decision
whether or not to skip the step.  This may result in one optimizer skipping the step
while the other one does not.  Since step skipping occurs rarely (every several hundred iterations)
this should not impede convergence.  If you observe poor convergence after adding gradient scaling
to a multiple-optimizer model, please file an issue.</p>
</div>
</div>
</div>


             </article>
             
            </div>
            <footer>
  
    <div class="rst-footer-buttons" role="navigation" aria-label="footer navigation">
      
        <a href="autograd.html" class="btn btn-neutral float-right" title="Autograd mechanics" accesskey="n" rel="next">Next <img src="../_static/images/chevron-right-orange.svg" class="next-page"></a>
      
      
        <a href="../index.html" class="btn btn-neutral" title="PyTorch documentation" accesskey="p" rel="prev"><img src="../_static/images/chevron-right-orange.svg" class="previous-page"> Previous</a>
      
    </div>
  

  

    <hr>

  

  <div role="contentinfo">
    <p>
        &copy; Copyright 2019, Torch Contributors.

    </p>
  </div>
    
      <div>
        Built with <a href="http://sphinx-doc.org/">Sphinx</a> using a <a href="https://github.com/rtfd/sphinx_rtd_theme">theme</a> provided by <a href="https://readthedocs.org">Read the Docs</a>.
      </div>
     

</footer>

          </div>
        </div>

        <div class="pytorch-content-right" id="pytorch-content-right">
          <div class="pytorch-right-menu" id="pytorch-right-menu">
            <div class="pytorch-side-scroll" id="pytorch-side-scroll-right">
              <ul>
<li><a class="reference internal" href="#">Automatic Mixed Precision examples</a><ul>
<li><a class="reference internal" href="#gradient-scaling">Gradient Scaling</a><ul>
<li><a class="reference internal" href="#typical-use">Typical Use</a></li>
<li><a class="reference internal" href="#working-with-unscaled-gradients">Working with Unscaled Gradients</a><ul>
<li><a class="reference internal" href="#gradient-clipping">Gradient clipping</a></li>
</ul>
</li>
<li><a class="reference internal" href="#working-with-scaled-gradients">Working with Scaled Gradients</a><ul>
<li><a class="reference internal" href="#gradient-penalty">Gradient penalty</a></li>
</ul>
</li>
<li><a class="reference internal" href="#working-with-multiple-losses-and-optimizers">Working with Multiple Losses and Optimizers</a></li>
</ul>
</li>
</ul>
</li>
</ul>

            </div>
          </div>
        </div>
      </section>
    </div>

  


  

     
       <script type="text/javascript" id="documentation_options" data-url_root="../" src="../_static/documentation_options.js"></script>
         <script src="../_static/jquery.js"></script>
         <script src="../_static/underscore.js"></script>
         <script src="../_static/doctools.js"></script>
         <script src="../_static/language_data.js"></script>
     

  

  <script type="text/javascript" src="../_static/js/vendor/popper.min.js"></script>
  <script type="text/javascript" src="../_static/js/vendor/bootstrap.min.js"></script>
  <script src="https://cdnjs.cloudflare.com/ajax/libs/list.js/1.5.0/list.min.js"></script>
  <script type="text/javascript" src="../_static/js/theme.js"></script>

  <script type="text/javascript">
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script>
 
<script>
  (function(i,s,o,g,r,a,m){i['GoogleAnalyticsObject']=r;i[r]=i[r]||function(){
  (i[r].q=i[r].q||[]).push(arguments)},i[r].l=1*new Date();a=s.createElement(o),
  m=s.getElementsByTagName(o)[0];a.async=1;a.src=g;m.parentNode.insertBefore(a,m)
  })(window,document,'script','https://www.google-analytics.com/analytics.js','ga');

  ga('create', 'UA-90545585-1', 'auto');
  ga('send', 'pageview');

</script>

<script async src="https://www.googletagmanager.com/gtag/js?id=UA-117752657-2"></script>

<script>
  window.dataLayer = window.dataLayer || [];

  function gtag(){dataLayer.push(arguments);}

  gtag('js', new Date());
  gtag('config', 'UA-117752657-2');
</script>

<img height="1" width="1" style="border-style:none;" alt="" src="https://www.googleadservices.com/pagead/conversion/795629140/?label=txkmCPmdtosBENSssfsC&amp;guid=ON&amp;script=0"/>


  <!-- Begin Footer -->

  <div class="container-fluid docs-tutorials-resources" id="docs-tutorials-resources">
    <div class="container">
      <div class="row">
        <div class="col-md-4 text-center">
          <h2>Docs</h2>
          <p>Access comprehensive developer documentation for PyTorch</p>
          <a class="with-right-arrow" href="https://pytorch.org/docs/stable/index.html">View Docs</a>
        </div>

        <div class="col-md-4 text-center">
          <h2>Tutorials</h2>
          <p>Get in-depth tutorials for beginners and advanced developers</p>
          <a class="with-right-arrow" href="https://pytorch.org/tutorials">View Tutorials</a>
        </div>

        <div class="col-md-4 text-center">
          <h2>Resources</h2>
          <p>Find development resources and get your questions answered</p>
          <a class="with-right-arrow" href="https://pytorch.org/resources">View Resources</a>
        </div>
      </div>
    </div>
  </div>

  <footer class="site-footer">
    <div class="container footer-container">
      <div class="footer-logo-wrapper">
        <a href="https://pytorch.org/" class="footer-logo"></a>
      </div>

      <div class="footer-links-wrapper">
        <div class="footer-links-col">
          <ul>
            <li class="list-title"><a href="https://pytorch.org/">PyTorch</a></li>
            <li><a href="https://pytorch.org/get-started">Get Started</a></li>
            <li><a href="https://pytorch.org/features">Features</a></li>
            <li><a href="https://pytorch.org/ecosystem">Ecosystem</a></li>
            <li><a href="https://pytorch.org/blog/">Blog</a></li>
            <li><a href="https://github.com/pytorch/pytorch/blob/master/CONTRIBUTING.md">Contributing</a></li>
          </ul>
        </div>

        <div class="footer-links-col">
          <ul>
            <li class="list-title"><a href="https://pytorch.org/resources">Resources</a></li>
            <li><a href="https://pytorch.org/tutorials">Tutorials</a></li>
            <li><a href="https://pytorch.org/docs/stable/index.html">Docs</a></li>
            <li><a href="https://discuss.pytorch.org" target="_blank">Discuss</a></li>
            <li><a href="https://github.com/pytorch/pytorch/issues" target="_blank">Github Issues</a></li>
            <li><a href="https://pytorch.org/assets/brand-guidelines/PyTorch-Brand-Guidelines.pdf" target="_blank">Brand Guidelines</a></li>
          </ul>
        </div>

        <div class="footer-links-col follow-us-col">
          <ul>
            <li class="list-title">Stay Connected</li>
            <li>
              <div id="mc_embed_signup">
                <form
                  action="https://twitter.us14.list-manage.com/subscribe/post?u=75419c71fe0a935e53dfa4a3f&id=91d0dccd39"
                  method="post"
                  id="mc-embedded-subscribe-form"
                  name="mc-embedded-subscribe-form"
                  class="email-subscribe-form validate"
                  target="_blank"
                  novalidate>
                  <div id="mc_embed_signup_scroll" class="email-subscribe-form-fields-wrapper">
                    <div class="mc-field-group">
                      <label for="mce-EMAIL" style="display:none;">Email Address</label>
                      <input type="email" value="" name="EMAIL" class="required email" id="mce-EMAIL" placeholder="Email Address">
                    </div>

                    <div id="mce-responses" class="clear">
                      <div class="response" id="mce-error-response" style="display:none"></div>
                      <div class="response" id="mce-success-response" style="display:none"></div>
                    </div>    <!-- real people should not fill this in and expect good things - do not remove this or risk form bot signups-->

                    <div style="position: absolute; left: -5000px;" aria-hidden="true"><input type="text" name="b_75419c71fe0a935e53dfa4a3f_91d0dccd39" tabindex="-1" value=""></div>

                    <div class="clear">
                      <input type="submit" value="" name="subscribe" id="mc-embedded-subscribe" class="button email-subscribe-button">
                    </div>
                  </div>
                </form>
              </div>

            </li>
          </ul>

          <div class="footer-social-icons">
            <a href="https://www.facebook.com/pytorch" target="_blank" class="facebook"></a>
            <a href="https://twitter.com/pytorch" target="_blank" class="twitter"></a>
            <a href="https://www.youtube.com/pytorch" target="_blank" class="youtube"></a>
          </div>
        </div>
      </div>
    </div>
  </footer>

  <div class="cookie-banner-wrapper">
  <div class="container">
    <p class="gdpr-notice">To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: <a href="https://www.facebook.com/policies/cookies/">Cookies Policy</a>.</p>
    <img class="close-button" src="../_static/images/pytorch-x.svg">
  </div>
</div>

  <!-- End Footer -->

  <!-- Begin Mobile Menu -->

  <div class="mobile-main-menu">
    <div class="container-fluid">
      <div class="container">
        <div class="mobile-main-menu-header-container">
          <a class="header-logo" href="https://pytorch.org/" aria-label="PyTorch"></a>
          <a class="main-menu-close-button" href="#" data-behavior="close-mobile-menu"></a>
        </div>
      </div>
    </div>

    <div class="mobile-main-menu-links-container">
      <div class="main-menu">
        <ul>
          <li>
            <a href="https://pytorch.org/get-started">Get Started</a>
          </li>

          <li>
            <a href="https://pytorch.org/features">Features</a>
          </li>

          <li>
            <a href="https://pytorch.org/ecosystem">Ecosystem</a>
          </li>

          <li>
            <a href="https://pytorch.org/mobile">Mobile</a>
          </li>

          <li>
            <a href="https://pytorch.org/hub">PyTorch Hub</a>
          </li>

          <li>
            <a href="https://pytorch.org/blog/">Blog</a>
          </li>

          <li>
            <a href="https://pytorch.org/tutorials">Tutorials</a>
          </li>

          <li class="active">
            <a href="https://pytorch.org/docs/stable/index.html">Docs</a>
          </li>

          <li>
            <a href="https://pytorch.org/resources">Resources</a>
          </li>

          <li>
            <a href="https://github.com/pytorch/pytorch">Github</a>
          </li>
        </ul>
      </div>
    </div>
  </div>

  <!-- End Mobile Menu -->

  <script type="text/javascript" src="../_static/js/vendor/anchor.min.js"></script>

  <script type="text/javascript">
    $(document).ready(function() {
      mobileMenu.bind();
      mobileTOC.bind();
      pytorchAnchors.bind();
      sideMenus.bind();
      scrollToAnchor.bind();
      highlightNavigation.bind();
      mainMenuDropdown.bind();
      filterTags.bind();

      // Remove any empty p tags that Sphinx adds
      $("[data-tags='null']").remove();

      // Add class to links that have code blocks, since we cannot create links in code blocks
      $("article.pytorch-article a span.pre").each(function(e) {
        $(this).closest("a").addClass("has-code");
      });
    })
  </script>
</body>
</html>